"""
This file provides classes and functions to match the peptides to the MS1-features
"""

import numpy 
import interval_tree
import data_handler


def match_peptides2feat(th_peptides, obs_features, \
        error_tolerance_list): 
    """
    Input: 
    th_peptides - List of theoretical peptides (data_handler.Peptide objects)
    obs_features - List of observed MS1-features (data_handler.Feature objects)
    error_tolerance_list - list including:
            mass tolerance, rt tolerance - if mass and rt errors treated 
                    separately
            avg mass error[ppm], sd of the mass errors, avg rt error, 
                    sd of the rt errors, radius of the accepted circle -
                    if the mass and rt errors are combined
    
    Output: The th_peptides will include the list of their matching features
    """   
    # set the appropriate functions        
    if len(error_tolerance_list) > 2:
        r = error_tolerance_list[4]
        fill_function = lambda peptides: fill_intervals_combined(\
                peptides, error_tolerance_list[0], \
                error_tolerance_list[1], r)
        matching_function = lambda f, p: are_matching_combined(f, p, \
                r*r, *error_tolerance_list[0:4])                
    else:
        fill_function = lambda peptides: fill_intervals_separate(\
                peptides, error_tolerance_list[0])
        matching_function = lambda f, p: are_matching_separate(f, p, \
                error_tolerance_list[1])

    # fill the mass intervals, and build an interval tree
    print "Building an interval tree, this may take a while..."
    fill_function(th_peptides)
    tree = interval_tree.IntervalTree(th_peptides)
    
    # try to find a matching peptide for each feature 
    print "Finding matching peptides for each feature, this may take a while..."
    for feat in obs_features:
        matching_mass = tree.search(feat.mass)
        matching_mass_rt = filter(lambda p: matching_function(feat, p), \
                matching_mass)
        if len(matching_mass_rt) > 0:
            matching_mass_rt = remove_duplicates(matching_mass_rt)
            for p in matching_mass_rt:
                p.matched_features.append(feat)
    print_statistics(th_peptides, obs_features)
    
    
def fill_intervals_combined(th_peptides, avg_mass_error, sd_mass_error, \
        radius):
    upper_offset = (radius*sd_mass_error + avg_mass_error)*1e-6 + 1.0
    lower_offset = 1.0 - (radius*sd_mass_error + avg_mass_error)*1e-6
    for i in xrange(len(th_peptides)):        
        th_mass = th_peptides[i].th_mass        
        upper_limit = th_mass * upper_offset
        lower_limit = th_mass * lower_offset
        th_peptides[i].mass_interval = (lower_limit, upper_limit) 


def are_matching_combined(feature, peptide, radius_sq, avg_mass_error, \
        sd_mass_error, avg_rt_error, sd_rt_error):            
    mass_error = numpy.abs(feature.mass-peptide.th_mass)/\
            peptide.th_mass*1e6
    rt_error = feature.rt-peptide.pred_rt
    norm_mass_error = (mass_error-avg_mass_error)/sd_mass_error
    norm_rt_error = (rt_error-avg_rt_error)/sd_rt_error
    if (norm_mass_error*norm_mass_error + norm_rt_error*norm_rt_error) <= \
            radius_sq:
        return True
    return False
  
  
def fill_intervals_separate(th_peptides, mass_tolerance):
    for i in xrange(len(th_peptides)):
        offset = float(mass_tolerance) * th_peptides[i].th_mass * 1e-6
        th_peptides[i].mass_interval = (th_peptides[i].th_mass - offset, \
                th_peptides[i].th_mass + offset)      
      

def are_matching_separate(feature, peptide, rt_tolerance_win):                    
    lower_limit = peptide.pred_rt - rt_tolerance_win
    upper_limit = peptide.pred_rt + rt_tolerance_win                
    if (feature.rt >= lower_limit and feature.rt <= upper_limit):
      return True
    return False


def remove_duplicates(list_):
    unique = []
    for item in list_:
        if not item in unique:  
            unique.append(item)
    return unique  


def print_statistics(th_peptides, obs_features):
    features_dict = {}
    pep_matching = []
    for p in th_peptides:
        if  len(p.matched_features) > 0:
            pep_matching.append(p.sequence)
            for feat in p.matched_features:
                if feat.id in features_dict:
                    features_dict[feat.id].append(p.sequence)
                else:
                    features_dict[feat.id] = [ p.sequence ]
    np = len(pep_matching)
    Np = len(th_peptides)
    print "%d/%d = %.1f%% peptides have at least one MS1-feature matching" %\
            (np, Np, float(np)/Np * 100)   
    
    nf = len(features_dict)
    Nf = len(obs_features)
    print "%d/%d = %.1f%% MS1-features have at least one peptide matching" %\
            (nf, Nf, float(nf)/Nf * 100)  

  
        
        


